package com.pearl.generator.service.impl;

import cn.hutool.core.util.ObjectUtil;
import com.baomidou.mybatisplus.annotation.DbType;
import com.baomidou.mybatisplus.core.toolkit.StringUtils;
import com.baomidou.mybatisplus.generator.config.DataSourceConfig;
import com.baomidou.mybatisplus.generator.config.IDbQuery;
import com.baomidou.mybatisplus.generator.config.StrategyConfig;
import com.baomidou.mybatisplus.generator.config.builder.ConfigBuilder;
import com.baomidou.mybatisplus.generator.config.builder.Entity;
import com.baomidou.mybatisplus.generator.config.po.TableField;
import com.baomidou.mybatisplus.generator.config.po.TableInfo;
import com.baomidou.mybatisplus.generator.config.querys.DbQueryDecorator;
import com.baomidou.mybatisplus.generator.config.rules.IColumnType;
import com.baomidou.mybatisplus.generator.jdbc.DatabaseMetaDataWrapper;
import com.pearl.generator.common.BusinessException;
import com.pearl.generator.common.Result;
import com.pearl.generator.pojo.dto.TableFieldDTO;
import com.pearl.generator.pojo.entity.DataSourceEntity;
import com.pearl.generator.pojo.mapstruct.TableInfoMapStructMapper;
import com.pearl.generator.pojo.request.TableInfoQuery;
import com.pearl.generator.pojo.vo.TableNameVO;
import com.pearl.generator.service.IDataSourceService;
import com.pearl.generator.service.TableInfoService;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Service;

import javax.annotation.Resource;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.util.*;
import java.util.stream.Collectors;

/**
 * @author TangDan
 * @version 1.0
 * @since 2022/8/4
 */
@Service
@Slf4j
public class TableInfoServiceImpl implements TableInfoService {

    @Resource
    IDataSourceService dataSourceService;

    @Resource
    TableInfoMapStructMapper tableInfoMapStructMapper;

    @Override
    public List<TableNameVO> selectAll(Long dataSourceId) {
        // 1. 根据ID 查询数据源
        DataSourceEntity sourceEntity = dataSourceService.getById(dataSourceId);
        if (ObjectUtil.isNull(sourceEntity)) {
            throw new IllegalArgumentException("未查询到当前数据源");
        }
        // 2. 数据源配置
        DataSourceConfig dataSourceConfig = new DataSourceConfig.Builder(sourceEntity.getUrl(), sourceEntity.getUsername(), sourceEntity.getPassword()).build();
        // 3. 策略配置
        StrategyConfig strategyConfig = new StrategyConfig.Builder().build();
        DbQueryDecorator dbQuery = new DbQueryDecorator(dataSourceConfig, strategyConfig);
        // 4. 执行查询
        List<TableNameVO> tables = new ArrayList<>();
        try {
            dbQuery.execute(dbQuery.tablesSql(), (result) -> {
                String tableName = result.getStringResult(dbQuery.tableName());
                if (StringUtils.isNotBlank(tableName)) {
                    TableNameVO tableNameVO = new TableNameVO();
                    tableNameVO.setTableName(tableName);
                    tableNameVO.setTableComment(result.getTableComment());
                    tables.add(tableNameVO);
                }
            });
        } catch (RuntimeException | SQLException runtimeException) {
            throw new BusinessException("查询失败，请检查数据源信息是否正确");
        } finally {
            dbQuery.closeConnection();
        }
        return tables;
    }

    @Override
    public List<TableFieldDTO> getOne(Long dataSourceId, String tableName) {
        // 1. 根据ID 查询数据源
        DataSourceEntity sourceEntity = dataSourceService.getById(dataSourceId);
        if (ObjectUtil.isNull(sourceEntity)) {
            throw new IllegalArgumentException("未查询到当前数据源");
        }
        // 2. 数据源配置
        DataSourceConfig dataSourceConfig = new DataSourceConfig.Builder(sourceEntity.getUrl(), sourceEntity.getUsername(), sourceEntity.getPassword()).build();
        // 3. 策略配置
        StrategyConfig strategyConfig = new StrategyConfig.Builder().addInclude(tableName).build();
        DbQueryDecorator dbQuery = new DbQueryDecorator(dataSourceConfig, strategyConfig);
        // 4. 查询数据
        return select(tableName, dataSourceConfig, dbQuery);
    }

    /**
     * 查询表结构信息
     *
     * @param tableName        表名
     * @param dataSourceConfig 数据源
     * @param dbQuery          查询对象
     * @return 表结构信息
     */
    private List<TableFieldDTO> select(String tableName, DataSourceConfig dataSourceConfig, DbQueryDecorator dbQuery) {
        List<TableFieldDTO> list =new ArrayList<>();
        DbType dbType = dataSourceConfig.getDbType();
        try {
            String tableFieldsSql = dbQuery.tableFieldsSql(tableName); // show full fields from `config_info_aggr`
            Set<String> h2PkColumns = new HashSet<>();
            if (DbType.H2 == dbType) {
                dbQuery.execute(String.format("select * from INFORMATION_SCHEMA.INDEXES WHERE TABLE_NAME = '%s'", tableName), (result) -> {
                    String primaryKey = result.getStringResult(dbQuery.fieldKey());
                    if (Boolean.parseBoolean(primaryKey)) {
                        h2PkColumns.add(result.getStringResult(dbQuery.fieldName()));
                    }
                });
            }
            dbQuery.execute(tableFieldsSql, (result) -> {
                TableFieldDTO tableFieldDTO = new TableFieldDTO();

                String columnName = result.getStringResult(dbQuery.fieldName()); /// 字段名
                tableFieldDTO.setColumnName(columnName);
                boolean isId = DbType.H2 == dbType ? h2PkColumns.contains(columnName) : result.isPrimaryKey(); // 是否主键
                tableFieldDTO.setKeyFlag(isId);
                if (isId) {
                    tableFieldDTO.setKeyIdentityFlag(dbQuery.isKeyIdentity(result.getResultSet())); // keyIdentityFlag 是否自动递增
                }
                tableFieldDTO.setType(result.getStringResult(dbQuery.fieldType()));// 类型
                tableFieldDTO.setComment(result.getFiledComment());// 字段注释
                Map<String, DatabaseMetaDataWrapper.ColumnsInfo> columnsMetaInfoMap = new HashMap<>();
                try {
                    columnsMetaInfoMap = (new DatabaseMetaDataWrapper(dbQuery.getConnection())).getColumnsInfo((String) null, dataSourceConfig.getSchemaName(), tableName);
                } catch (SQLException e) {
                    e.printStackTrace();
                    throw new BusinessException("查询失败，请检查数据源信息是否正确");
                }
                tableFieldDTO.setMetaInfo(new TableField.MetaInfo(columnsMetaInfoMap.get(columnName.toLowerCase())));
                list.add(tableFieldDTO);
            });
            return list;
        } catch (SQLException e) {
            e.printStackTrace();
            throw new BusinessException("查询失败，请检查数据源信息是否正确");
        }
    }
}